package cn.phecda.boot.tenant.core.spring;

import cn.phecda.boot.common.lang.constants.DefaultConsts;
import cn.phecda.boot.common.lang.model.TenantContext;
import cn.phecda.boot.tenant.core.service.TenantDataSourceService;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.lang.NonNull;

import javax.sql.DataSource;

/**
 * <p> 租户数据源路由 </p>
 *
 * @author : wanpinwei
 */
@Slf4j
@RequiredArgsConstructor
public class TenantRoutingDataSource extends AbstractRoutingDataSource {
    private final TenantDataSourceService tenantDataSourceService;
    private final DataSource dataSource;

    @PreDestroy
    public void destroy() {
        tenantDataSourceService.destory();
    }

    @Override
    @NonNull
    protected DataSource determineTargetDataSource() {
        String lookupKey = (String)this.determineCurrentLookupKey();
        if (DefaultConsts.TENTANT_CODE.equalsIgnoreCase(lookupKey)) {
            // 使用主数据源
            return dataSource;
        }

        return tenantDataSourceService.getDataSource(lookupKey);
    }


    @Override
    protected Object determineCurrentLookupKey() {
        return TenantContext.get();
    }

    @Override
    public void initialize() {
        this.setDefaultTargetDataSource(dataSource);
    }
}
